import math
import random # Python's RNG is very good IIRC, so no need to reinvent the wheel

class DiamondSquare:
	def __init__(self, minimalSize):
		n = math.log( minimalSize - 1 ) / math.log( 2 ) # size must be power of two plus one
		n = int( 0.5 + math.ceil( n ) )
		self.size = 2 ** n + 1 
		self.scale = self.size // 2 # "radius"
		self.data = []
		for i in range(self.size):
			line = []
			for j in range(self.size):
				line.append( None )
			self.data.append( line )
		self.squareQ = []
		self.diamondQ = []
		self.data[0][0] = 0.0
		self.data[self.size-1][0] = 0.0
		self.data[0][self.size-1] = 0.0
		self.data[self.size-1][self.size-1] = 0.0
	def setCenter(self, radius, value):
		for x in range(self.size):
			for y in range(self.size):
				dx, dy = x - self.size // 2, y - self.size // 2
				if dx*dx + dy*dy < radius*radius:
					self.data[x][y] = value
	def zeroEdges(self):
		for x in range(self.size):
			for y in range(self.size):
				if x == 0 or (self.size == (x+1)) or y == 0 or (self.size == (y+1)):
					self.data[x][y] = 0
	def setOutsideRegion(self, region, value):
		x0, y0, width, height = region
		for x in range(self.size):
			for y in range(self.size):
				if x >= x0 and x < (x0+width) and y >= y0 and y < (y0+height): continue
				self.data[x][y] = value
	def absolutize(self):
		for x in range(self.size):
			for y in range(self.size):
				self.data[x][y] = abs( self.data[x][y] )
	def normalize(self, region = None):
		mx, mn = None, None
		if region:
			x0, y0, width, height = region
		else:
			x0, y0, width, height = 0,0,self.size, self.size
		for x in range(width):
			for y in range(height):
				if mx == None or self.data[x+x0][y+y0] > mx:
					mx = self.data[x+x0][y+y0]
				if mn == None or self.data[x+x0][y+y0] < mn:
					mn = self.data[x+x0][y+y0]
		for x in range(width):
			for y in range(height):
				self.data[x+x0][y+y0] = ( self.data[x+x0][y+y0] - mn ) / float( mx - mn )
	def generate(self):
		self.done = set()
		self.diamond( self.size // 2, self.size // 2 )
		self.done.add( (self.size // 2, self.size // 2 ) )
		while self.squareQ:
			while self.squareQ:
				x, y = self.squareQ.pop(0)
				self.square( x, y )
			self.scale /= 2
			while self.diamondQ:
				x, y = self.diamondQ.pop(0)
				self.diamond( x, y )
	def perturbation(self):
		mopo = random.random() * 2 - 1
		return mopo * self.scale
	def diamond( self, x, y ):
		if self.data[x][y] == None:
			left, right = x - self.scale, x + self.scale
			top, bottom = y - self.scale, y + self.scale
			values = []
			for tx, ty in [ (left,top), (right,top), (left,bottom), (right, bottom) ]:
				if tx < 0 or tx >= self.size: continue
				if ty < 0 or ty >= self.size: continue
				values.append( self.data[tx][ty] )
			self.data[x][y] = sum( values ) / float( len( values ) ) + self.perturbation()
		for dx, dy in [ (dx,dy) for dx in (-1,0,1) for dy in (-1,0,1) if (dx or dy) and not (dx and dy) ]:
			tx, ty = (x+self.scale*dx, y + self.scale*dy )
			if not (tx,ty) in self.done:
				self.squareQ.append( (x+self.scale*dx, y + self.scale*dy ) )
				self.done.add( (tx,ty) )
	def square( self, x, y ):
		if self.data[x][y] == None:
			values = []
			for dx, dy in [ (dx,dy) for dx in (-1,0,1) for dy in (-1,0,1) if (dx or dy) and not (dx and dy) ]:
				tx, ty = x + dx * self.scale, y + dy * self.scale
				if tx < 0 or tx >= self.size: continue
				if ty < 0 or ty >= self.size: continue
				values.append( self.data[tx][ty] )
			self.data[x][y] = sum( values ) / float( len( values ) ) + self.perturbation()
		scale = self.scale // 2
		if scale:
			for dx, dy in [ (dx,dy) for dx in (-1,0,1) for dy in (-1,0,1) if dx and dy ]:
				tx, ty = x + dx * scale , y + dy * scale
				if tx < 0 or tx >= self.size: continue
				if ty < 0 or ty >= self.size: continue
				if not (tx,ty) in self.done:
					self.diamondQ.append( (tx,ty) )
					self.done.add( (tx,ty) )
	def display( self ):
		import Image
		img = Image.new( "RGB", (self.size,self.size) )
		mx, mn = None, None
		for x in range(self.size):
			for y in range(self.size):
				v = self.data[x][y]
				if mx == None or v > mx: mx = v
				if mn == None or v < mn: mn = v
		pic = img.load()
		for x in range(self.size):
			for y in range(self.size):
				v = int( 0.5 + 255.0 * (self.data[x][y] - mn) / float(mx-mn) )
				pic[ x, y ] = v, v, v
		img.show()

if __name__ == '__main__':
	import Image
	import time
	import psyco
	psyco.full()
	t0 = time.time()
	sz = 160
	red = DiamondSquare( sz )
	blue = DiamondSquare( sz )
	green = DiamondSquare( sz )
	for o in (red,blue,green):
		o.zeroEdges()
		o.setOutsideRegion( (0,0,o.size,o.size//3) , 0 )
		o.generate()
		o.absolutize()
		o.normalize()
	img = Image.new( "RGB", (red.size,red.size) )
	pic = img.load()
	for x in range(red.size):
		for y in range(red.size):
			r = int( red.data[x][y] * 255.0 + 0.5 )
			g = int( green.data[x][y] * 255.0 + 0.5 )
			b = int( blue.data[x][y] * 255.0 + 0.5 )
			pic[ x, y ] = r, g, b
	print time.time() - t0
	img.show()
	img.save( "diamond_square.png" )
